
from ast import arg
from PIL import Image
import glob 
import os 
from os.path import join, split,exists
from tqdm import tqdm
from img_utils import *



import pdb

import argparse


def get_preds(src,method_suffix = 'step0'):
    references = sorted(glob.glob(src+'/*ref.png'))
    gts = sorted(glob.glob(src+'/*gt.png'))
    prediction = sorted(glob.glob(src+'/*%s.png'%(method_suffix)))
    errors = sorted(glob.glob(src+'/*error.png'))

    return references,gts,prediction,errors




def compare_two_src():

    src1= "/baai-cwm-1/baai_cwm_ml/cwm/shaocong.xu/exp/StableNormal_training/eval_results/DIODE_normal/yoso-normal-v-1-0_xt-2025-02-09#07:26:45.031091-seed352339-resolution_1024"

    src2= "/baai-cwm-1/baai_cwm_ml/cwm/shaocong.xu/exp/StableNormal_training/eval_results/DIODE_normal/yoso-normal-v-1-0_xt-2025-02-09#06:24:41.593566-seed337891-resolution_1024"


    references,gts,prediction,errors = get_preds(src1)
    _,_, prediction2, errors2 = get_preds(src2)


    length = len(references)
    # assert len(references) == len(prediction)

    print(length)

    output_path = '/baai-cwm-1/baai_cwm_ml/cwm/shaocong.xu/exp/StableNormal_training/eval_results/DIODE_normal/outputs'
    if not exists(output_path):
        os.makedirs(output_path)



    pil_imgs = []
    for i in  tqdm(range(length)):

        ref = references[i]
        gt = gts[i]
        pred = prediction[i]
        err = errors[i]

        pred2 = prediction2[i]
        err2 = errors2[i]


        # pil_imgs.append(Image.open(imgs[i]))
        
        output_name = join(output_path,pred.split('/')[-1])
        if True or not exists(output_name):
            if exists(ref) and exists(pred):
                img1 = concat_images([ref,gt,pred,err],direction='horizontal')
                img2 = concat_images([ref,gt,pred2,err2],direction = 'horizontal')
                final = concat_images([img1,img2],output_name,'vertical')


        if i == 5:
            break




if __name__ == "__main__":

    
    parser = argparse.ArgumentParser()
    parser.add_argument('--root_path', type=str, default="./logs/outputs")
    parser.add_argument('--output_path', type=str, default="./logs")
    parser.add_argument('--vis_num', type=int, default=10)
    parser.add_argument('--method_suffix', type=str, default='step0')
    args = parser.parse_args()


    src1= args.root_path

    references,gts,prediction,errors = get_preds(src1,method_suffix = args.method_suffix)
    
    
    length = len(references)

    print(length)
    output_path = args.output_path


    if not exists(output_path):
        os.makedirs(output_path)

    pil_imgs = []

    
    

    sampled_idx = np.random.choice(list(range(length)), args.vis_num, replace =False)

    print('sample idx number : %d'%(len(sampled_idx)))

    
    # for i in  tqdm(range(length)):
    for i in  tqdm(sampled_idx):
        
        ref = references[i]
        gt = gts[i]
        pred = prediction[i]
        err = errors[i]

        output_name = join(output_path,pred.split('/')[-1])
        if True or not exists(output_name):
            if exists(ref) and exists(pred):
                img1 = concat_images([ref,gt,pred,err],output_name,direction='horizontal')

        
